import wandb
import os
from pathlib import Path
import torch
from DeepPPG.config.base import BaseConfig


class PredConfig(BaseConfig):
    project_name = "AFPredictionDebug"
    seed = 42
    resume = False  # do not resume training
    sweep = False  # whether to sweep hyperparameters, if True, will ignore tunable hyperparameters below
    # W&B configuration
    wandb_config = None
    metric = "f1_score"
    device = "cuda" if torch.cuda.is_available() else "cpu"

    training_type = 'seq2one'
    label_col = 'AF_1year'
    data_dir = 'g:\\/NEXUS/ECG_files/'
    train_data_dir = 'DeepECG/dataset/alpine/train_1_year.csv'
    val_data_dir = 'DeepECG/dataset/alpine/val_1_year.csv'
    test_data_dir = 'DeepECG/dataset/alpine/test_1_year.csv'
    lead_type = ['I','II','III','V1','V2','V3','V4','V5','V6','aVR','aVL','aVF']
    sig_length = 2048
    bandpass_lower = 0.1
    bandpass_higher = 100
    median_filter_size = 5
    model_save_dir = None
    log_dir = None

    dataloader_workers = 8
    print_time = False
    num_epochs = 80
    early_stopping = 10

    # general hyperparameters
    verbose = True
    batch_size = 1000
    eval_batch_size = 500
    lr = 1e-3
    dropout = 0.
    weight_decay = 0.
    hidden_size = 256

class ResNetConfig(PredConfig):
    model_type = 'MetaResNet'  # Model type: 
    do_meta = True
    if_t_dist = True
    # Model configuration
    grad_bound = 5
    al = 0.01
    la = 1e-3
    num_classes = 1
    input_channels = 12
    output_channels = 256
    middle_channels = 64
    criterion = {"criterion":"FocalBCE", "criterion_hp":{"alpha": 0.5, "gamma": 2.0}}
    # criterion = {"criterion": "BCE", "criterion_hp": None}
    log_dir = f"{PredConfig.training_type}_{model_type}_log"
    model_save_dir = Path(log_dir) / 'models'
    